Background#
We are going to demonstrate the use of the split-operator method to solve the time-dependent Schroedinger equation in 1D,
\[ i\hbar \frac{\partial}{\partial t} \Psi(x,t) = \hat{H} \Psi(x,t), \]
where
\[ \begin{align}\begin{aligned} \hat{H} = -\frac{\hbar^2}{2m} \frac{d^2}{dx^2} + V(x) \equiv \hat{K} + \hat{V} $$.\\The idea behind the split operator method is to approximate the time evolution operator $\hat{U}(t,t_0)$ for a very small time-increment $\Delta t$ as follows:\\$$ \hat{U}(t+\Delta t, t) = {\rm exp}\left( -\frac{i \Delta t}{2\hbar} \hat{V}\right)
{\rm exp}\left( -\frac{i\Delta t}{\hbar} \hat{K} \right)
{\rm exp}\left( - \frac{i \Delta t}{2\hbar} \hat{V} \right) $$,
such that\\$$ \Psi(x, t+\Delta t) \approx \hat{U}(t+\Delta t, t) \Psi(x,t). \end{aligned}\end{align} \]
In this approach, we apply a half time-step update to the wavefunction through the kinetic energy operator, a full time-step update through the potential energy operator, then a second half time-step update through the kinetic energy operator.
That looks great, but how the heck do we exponentiate \(\hat{K}\) with it’s second derivatives? IDK, but actually we don’t have to. The kinetic energy operator Fourier transforms to a quadratic function in momentum space:
\[ -\frac{\hbar^2}{2m}\frac{d^2}{dx^2} \xrightarrow{\mathcal{F}} \frac{ p^2}{2m}, \]
and this can be easily exponentiated. We will denote the kinetic energy operator in momentum space as \( \hat{K}_p = \frac{p^2}{2m} \), and we can therefore write the following:
\[ {\rm exp}\left( -\frac{i \Delta t}{\hbar} \hat{K}\right) \Psi(x,t) = \mathcal{F}^{-1}\left( {\rm exp}\left( -\frac{i \Delta t}{\hbar} \hat{K}_p \right) \mathcal{F} \left( \Psi(x,t) \right) \right). \]
For our purposes, this is actually awesome because we can take the Fourier and inverse Fourier transform of things with like a single line of python code. The only mild additional discomfort we have to endure is having to carry an additional grid around - we will need a spatial grid to evaluate \(\Psi(x,t)\) and \(\hat{V}\) on, but we will need a momentum grid to evaluate \(\hat{K}_p\) on.
We can define the points on our spatial grid as follows:
\[ \begin{align}\begin{aligned} x_j = x_{min} + j \Delta x $$ for $ j = 0, 1, ..., N-1 $ where $N$ is the number of gridpoints and\\$$ \Delta x = \frac{x_{max} - x_{min}}{N-1}$$. \\Our momentum grid will be defined as follows:\\$$ p_j = j \frac{2 \pi \hbar}{N \Delta x}, \end{aligned}\end{align} \]
where \(j = -\frac{N}{2}, ..., \frac{N}{2}. \)
Outline of our code#
Add attributes to our class for \(x_{min}\), \(x_{max}\), \(N\) and create grids
Add attributes to our class for other parts of the simulation, including
time step \(\Delta t\) (in atomic units)
mass of particle \(m\) (in atomic units)
value of \(\hbar\) (in atomic units)
Define initial wavefunction $\(\Psi(x,t_0) = \frac{1}{\sigma \sqrt{2\pi}} {\rm exp}\left(\frac{-1}{2}\left(\frac{x-x_0}{\sigma}\right)^2\right) {\rm exp}\left(i k_0 x \right)\)$
\(k_0 \approx \frac{m}{ 2}\)
\(x_0 \approx \frac{x_{max}}{2}\)
\(\sigma \approx \frac{x_{max}-x_{min}}{30} \)
Define \(\hat{V}\)… could be
Nothing
Square barrier
Square well
Gaussian barrier?
Harmonic well?
Define \(\hat{K}_p\) operator using \(\frac{{\bf p}^2}{2m}\)
Form exponentials of the \(\hat{V}\) and \(\hat{K}_p\) operators
\({\bf K} = {\rm exp}\left( -\frac{i \Delta t}{\hbar} \hat{K}_p \right) \)
\({\bf V} = {\rm exp}\left( -\frac{i \Delta t}{2\hbar} \hat{V} \right) \)
Implement split operator method:
\(\Psi^{(1)}(x,t) = {\bf V} \: \Psi(x,t) \)
\(\Psi^{(1)}(p,t) = \mathcal{F} \left( \Psi^{(1)}(x,t)\right)\)
\(\Psi^{(2)}(p,t) = {\bf K} \: \Psi^{(1)}(p,t)\)
\(\Psi^{(2)}(x,t) = \mathcal{F}^{-1} \left( \Psi^{(2)}(p,t)\right) \)
\(\Psi(x,t+\Delta t) = {\bf V} \: \Psi^{(2)}(x,t) \)
from matplotlib import animation, rc, pyplot as plt
from IPython.display import HTML
import numpy as np
class Quantum:
def __init__(self, args):
if 'box_length' in args:
self.L = args['box_length']
else:
self.L = 10.0
if 'time_step' in args:
self.dt = args['time_step']
else:
self.dt = 0.1
if 'grid_points' in args:
self.grid_points =args['grid_points']
else:
self.grid_points = 256
if 'mass' in args:
self.m = args['mass']
else:
self.m = 1
''' some constants and parameters in atomic units '''
# reduced planck's constant
self.hbar = 1
# reduced mass for PIR system or a very light harmonic oscillator!
self.mu = 1
# very small spring constant!
self.k = 1
# position-space spread for a gaussian wavepacket
self.sigma = self.L / 5
# central position value for a gaussian wavepacket
self.x0 = self.L / 2
# central momentum for a gaussian wavepacket
self.k0 = 0.4
self.x = np.linspace(0, self.L, self.grid_points)
self.dx = self.x[1]-self.x[0]
res = self.grid_points
self.dk = 2 * np.pi / (res * self.dx)
self.k = np.concatenate((np.arange(0, res / 2),
np.arange(-res / 2, 0))) * self.dk
self.V = np.zeros_like(self.x)
#self.V = 0.5 * (self.x - self.voffset) ** 2
self.Psi = np.sqrt(2/self.L) * np.sin((np.pi * self.x/self.L), dtype=complex)
self.K = np.exp(-1/(2*self.m) * (self.k ** 2) * self.dt * 1j)
self.R = np.exp(-0.5 * self.V * self.dt * 1j)
def build_operator(self):
self.R = np.exp(-0.5 * self.V * self.dt * 1j)
def split_op(self):
# Half-step in real space
self.Psi *= self.R
# FFT to momentum space
self.Psi = np.fft.fft(self.Psi)
# Full step in momentum space
self.Psi *= self.K
# iFFT back
self.Psi = np.fft.ifft(self.Psi)
# Final half-step in real space
self.Psi *= self.R
params = {'box_length': 500, 'v_offset': 0, 'grid_points': 1000, 'wfc_offset': 0, 'system': 'pib', 'mass': 1.0,
'time_step': 1.0}
wf = Quantum(params)
ci = 0+1j
x0 = 200
sig = 15
k0 = 0.4*1
### T1 will be the prefactor that is 1/(sigma * sqrt(2 * pi))
T1 = 1/(sig * np.sqrt(2 * np.pi))
### T2 will be the Gaussian function, exp(-0.5 * ((x-x0)/sigma)^2)
T2 = np.exp(-0.5 * ((wf.x-x0)/sig)**2)
### T3 will be the complex exponential (aka the plane wave!)
T3 = np.exp(ci * k0 * wf.x)
wf.Psi = T1 * T2 * T3
### put in a square barrier between x=600 and x=700
wf.V[600:700] = +0.025
wf.build_operator()
# First set up the figure, the axis, and the plot element we want to animate
fig, ax = plt.subplots()
plt.close()
### parameters for plot
ax.set_xlim(( 0, 500))
ax.set_ylim((-0.003, 0.003))
line, = ax.plot([], [], lw=2)
lineV, = ax.plot([], [], lw=2)
# initialization function: plot the background of each frame
def init():
line.set_data([], [])
lineV.set_data([], [])
return (line, lineV,)
#def build_triangle_potential(leftpoint, midpoint, rightpoint):
def animate(i):
wf.split_op()
#print(i,wf.Psi[2000])
line.set_data(wf.x, np.real(np.conj(wf.Psi)*wf.Psi))
lineV.set_data(wf.x, wf.V)
return(line, lineV,)
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=400, interval=10, blit=True)
# Note: below is the part which makes it work on Colab
rc('animation', html='jshtml')
anim
# (Optional) Utility: infer a barrier right edge from a potential array
import numpy as np
def infer_barrier_right_edge(x, V, frac=0.1):
Vmax = np.max(V)
thresh = frac * Vmax if Vmax > 0 else 0.0
mask = V > thresh
if not np.any(mask):
return float(x[int(0.7*len(x))])
idx = np.where(mask)[0]
return float(x[idx.max()])
# === Feature 1: Interactive potential builder (rectangle/triangle/trapezoid) ===
import numpy as np
import ipywidgets as w
import matplotlib.pyplot as plt
from IPython.display import display, clear_output
def rectangular_potential(x, x_left, width, height):
x_right = x_left + width
return np.where((x >= x_left) & (x <= x_right), height, 0.0)
def triangular_potential(x, x_left, width, height, apex_pos=0.5):
x_right = x_left + width
V = np.zeros_like(x, dtype=float)
apex_x = x_left + apex_pos * width
left_mask = (x >= x_left) & (x <= apex_x)
if np.any(left_mask):
V[left_mask] = height * (x[left_mask] - x_left) / max(apex_x - x_left, 1e-15)
right_mask = (x > apex_x) & (x <= x_right)
if np.any(right_mask):
V[right_mask] = height * (x_right - x[right_mask]) / max(x_right - apex_x, 1e-15)
return V
def trapezoid_potential(x, x_left, width, height_left, height_right, ramp_frac=0.2):
w_ramp = max(ramp_frac * width, 1e-15)
x_rise_end = x_left + w_ramp
x_fall_start = x_left + width - w_ramp
x_right = x_left + width
V = np.zeros_like(x, dtype=float)
mask_rise = (x >= x_left) & (x < x_rise_end)
if np.any(mask_rise):
V[mask_rise] = height_left * (x[mask_rise] - x_left) / w_ramp
mask_mid = (x >= x_rise_end) & (x <= x_fall_start)
if np.any(mask_mid):
mid_len = max(x_fall_start - x_rise_end, 1e-15)
V[mask_mid] = height_left + (height_right - height_left) * (x[mask_mid] - x_rise_end) / mid_len
mask_fall = (x > x_fall_start) & (x <= x_right)
if np.any(mask_fall):
V[mask_fall] = height_right * (x_right - x[mask_fall]) / w_ramp
return V
def build_potential(x, shape, params):
if shape == "Rectangle":
return rectangular_potential(x, params["x_left"], params["width"], params["height"])
elif shape == "Triangle":
return triangular_potential(x, params["x_left"], params["width"], params["height"], params["apex_pos"])
elif shape == "Trapezoid":
return trapezoid_potential(x, params["x_left"], params["width"], params["height_left"], params["height_right"], params["ramp_frac"])
else:
return np.zeros_like(x)
shape_dd = w.Dropdown(options=["Rectangle", "Triangle", "Trapezoid"], value="Rectangle", description="Shape:")
x_left_sl = w.FloatSlider(value=float(x.min()) + 0.2*(x.max()-x.min()), min=float(x.min()), max=float(x.max()), step=float((x[1]-x[0])), description="x_left:", continuous_update=False)
width_sl = w.FloatSlider(value=0.2*(x.max()-x.min()), min=float((x[1]-x[0])*2), max=float(x.max()-x.min()), step=float((x[1]-x[0])), description="Width:", continuous_update=False)
height_sl = w.FloatSlider(value=1.0, min=0.0, max=10.0, step=0.01, description="Height:", continuous_update=False)
apex_pos_sl = w.FloatSlider(value=0.5, min=0.0, max=1.0, step=0.01, description="Apex pos:", continuous_update=False)
height_l_sl = w.FloatSlider(value=1.0, min=0.0, max=10.0, step=0.01, description="Left H:", continuous_update=False)
height_r_sl = w.FloatSlider(value=0.5, min=0.0, max=10.0, step=0.01, description="Right H:", continuous_update=False)
ramp_frac_sl = w.FloatSlider(value=0.2, min=0.0, max=0.45, step=0.01, description="Ramp frac:", continuous_update=False)
pot_out = w.Output()
V_current = {"V": np.zeros_like(x)}
def update_plot(*_):
with pot_out:
clear_output(wait=True)
shape = shape_dd.value
params = dict(
x_left=x_left_sl.value,
width=width_sl.value,
height=height_sl.value,
apex_pos=apex_pos_sl.value,
height_left=height_l_sl.value,
height_right=height_r_sl.value,
ramp_frac=ramp_frac_sl.value
)
V = build_potential(x, shape, params)
V_current["V"] = V
fig, ax = plt.subplots(figsize=(6,3))
ax.plot(x, V)
ax.set_title(f"{shape} potential")
ax.set_xlabel("x")
ax.set_ylabel("V(x)")
ax.grid(True, alpha=0.3)
plt.show()
def _toggle_controls(*_):
shape = shape_dd.value
height_sl.layout.display = "flex" if shape == "Rectangle" else "none"
apex_pos_sl.layout.display = "flex" if shape == "Triangle" else "none"
height_l_sl.layout.display = "flex" if shape == "Trapezoid" else "none"
height_r_sl.layout.display = "flex" if shape == "Trapezoid" else "none"
ramp_frac_sl.layout.display= "flex" if shape == "Trapezoid" else "none"
update_plot()
for wid in [shape_dd, x_left_sl, width_sl, height_sl, apex_pos_sl, height_l_sl, height_r_sl, ramp_frac_sl]:
wid.observe(update_plot, names="value")
shape_dd.observe(_toggle_controls, names="value")
_toggle_controls()
controls = w.VBox([shape_dd, x_left_sl, width_sl, height_sl, apex_pos_sl, height_l_sl, height_r_sl, ramp_frac_sl])
ui = w.HBox([controls, pot_out])
display(ui)
# To use in your solver: replace your V with V_current["V"] before the time loop.
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[5], line 52
49 return np.zeros_like(x)
51 shape_dd = w.Dropdown(options=["Rectangle", "Triangle", "Trapezoid"], value="Rectangle", description="Shape:")
---> 52 x_left_sl = w.FloatSlider(value=float(x.min()) + 0.2*(x.max()-x.min()), min=float(x.min()), max=float(x.max()), step=float((x[1]-x[0])), description="x_left:", continuous_update=False)
53 width_sl = w.FloatSlider(value=0.2*(x.max()-x.min()), min=float((x[1]-x[0])*2), max=float(x.max()-x.min()), step=float((x[1]-x[0])), description="Width:", continuous_update=False)
54 height_sl = w.FloatSlider(value=1.0, min=0.0, max=10.0, step=0.01, description="Height:", continuous_update=False)
NameError: name 'x' is not defined
# === Feature 2: Transmitted probability via integrated probability current ===
import numpy as np
def probability_current_1d(psi, dx, hbar=1.0, m=1.0):
dpsi_dx = np.zeros_like(psi, dtype=complex)
dpsi_dx[1:-1] = (psi[2:] - psi[:-2]) / (2*dx)
dpsi_dx[0] = (psi[1] - psi[0]) / (dx)
dpsi_dx[-1] = (psi[-1] - psi[-2]) / (dx)
j = (hbar/m) * np.imag(np.conj(psi) * dpsi_dx)
return j
class FluxDetector:
def __init__(self, x, x_det_index, dt, hbar=1.0, m=1.0, smooth=0):
self.x = x
self.idx = int(x_det_index)
self.dt = float(dt)
self.hbar = float(hbar)
self.m = float(m)
self.accum = 0.0
self._buffer = []
self._smooth = int(smooth)
def update(self, psi):
dx = self.x[1] - self.x[0]
j = probability_current_1d(psi, dx, self.hbar, self.m)
val = float(j[self.idx])
if self._smooth > 0:
self._buffer.append(val)
if len(self._buffer) > self._smooth:
self._buffer.pop(0)
val = sum(self._buffer)/len(self._buffer)
self.accum += val * self.dt
@property
def P_transmitted(self):
return float(self.accum)
# --- Quick start ---
# x_d = infer_barrier_right_edge(x, V) + 3*(x[1]-x[0]) # small buffer
# x_det_index = int(np.argmin(np.abs(x - x_d)))
# det = FluxDetector(x, x_det_index, dt, hbar=hbar, m=m, smooth=2)
# In your time loop (after updating K each step): det.update(K)
# After the loop: print("Transmitted probability (flux):", det.P_transmitted)